import datetime
import pandas as pd
import tushare as ts
import matplotlib.pyplot as plt
from matplotlib.widgets import MultiCursor
from pandas.plotting import register_matplotlib_converters

import QUANTAXIS as QA
from QAStrategy.QAStrategy.classic_strategy.turtle_strategy import TurtleTrade


register_matplotlib_converters()


def main(codes, start, end):
    # print(codes.head())
    data = QA.QA_fetch_stock_day_adv(codes['code'].to_list(), start, end)
    data = data.to_qfq()
    # ind = data.add_func(QA.QA_indicator_THS_ZLXC2)
    ind = data.add_func(QA.QA_indicator_THS_3P)
    # ind = data.groupby(level=1, sort=False).apply(QA.QA_indicator_Turtle,
    #                                                           N_UP=55,
    #                                                           N_DN=20,
    #                                                           N_S_MA=10,
    #                                                           N_L_MA=60,
    #                                                           N_ATR=20)
    for _, code in codes.iterrows():
        df = ind.xs(code['code'], level='code', axis=0, drop_level=True)
        # df.to_csv('601933_zlxc.csv')
        ind_last_last_bar = df.iloc[-3]
        ind_last_bar = df.iloc[-2]
        ind_current_bar = df.iloc[-1]
        # if ind_current_bar.VAR1 > 50 or ind_last_bar.VAR1 > 50 or ind_last_last_bar.VAR1 > 50:
        # if ind_current_bar.VAR5 > 50 or ind_last_bar.VAR5 > 50:
        if ind_current_bar.B > 0 or ind_last_bar.B > 0:
            # and ind_current_bar.name[0].year == 2020 \
            # and ind_current_bar.name[0].month == 12 \
            # and ind_current_bar.name[0].day >= 1:
            print('code：{}---{}---时间{}---CLOSE{}'.format(
                code['code'], code['name'], ind_current_bar.name, ind_current_bar.CLOSE))


            figure = plt.figure()
            plt.title('{}---'.format(code['code']))
            axes1 = figure.add_subplot(411)
            axes2 = figure.add_subplot(412)
            axes3 = figure.add_subplot(413)
            axes4 = figure.add_subplot(414)
            axes1.plot(df.index, df['CLOSE'])
            axes2.plot(df.index, df['B'])
            axes2.plot(df.index, df['S'])
            # axes2.plot(df.index, df['VAR5'])
            axes3.plot(df.index, df['VAR1'])
            axes3.plot(df.index, df['VAR2'])
            axes3.plot(df.index, df['VAR3'])
            axes4.plot(df.index, df['X'])
            axes4.plot(df.index, df['Y'])

            left, right = axes1.get_xlim()
            # axes1.hlines(y=box.high, xmin=left, xmax=right, linestyles='dashed')
            multi = MultiCursor(figure.canvas, (axes1, axes2, axes3, axes4), color='r', lw=1)

            plt.show()

    # data.to_csv('002581.csv')
    # print(data.index)


if __name__ == '__main__':
    ts_token = '17056d23a59ab71cb979c6a30185e092aba605c4544dac900a3eb7f8'
    ts.set_token(ts_token)
    pro = ts.pro_api()
    data = pro.stock_basic(exchange='', list_status='L', fields='symbol,name,area,industry,list_date')
    data = data[data.symbol.str.startswith('00') | data.symbol.str.startswith('60')]
    # data = data[data.symbol.str.startswith('601933')]
    data = data[data.list_date < '20160101']
    all = data.rename(columns={'symbol': 'code'})
    codes = all[~all.name.str.contains('ST')]
    # codes = codes.iloc[0: 1]
    # hs300 = pd.read_csv('D:\\PythonPro\\QUANTAXIS\\AI\\hs300.csv')
    # hs300 = ts.get_hs300s()
    # codes = codes['code'].to_list()
    # codes = ['002027']
    # print(codes)
    start = '2020-01-01'
    end = '2020-12-31'
    main(codes, start, end)
